#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import time
import random
import pymysql
import logging
from datetime import datetime, timedelta
from tqdm import tqdm

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'port': 3306,
    'user': 'root',
    'password': '13177936413qq',
    'db': 'codoon_data',
    'charset': 'utf8mb4'
}

class CodoonBatchGenerator:
    """咕咚运动大批量数据生成器"""
    
    def __init__(self, batch_size=500):
        """初始化数据生成器"""
        self.conn = None
        self.cursor = None
        self.batch_size = batch_size  # 批量插入大小
        
        # 用于生成更真实数据的预设值
        self.realistic_user_names = [
            "跑步达人", "马拉松爱好者", "晨跑族", "夜跑狂人", "越野跑者", 
            "健身达人", "运动博主", "铁人三项", "骑行达人", "健康生活家",
            "山地车手", "徒步爱好者", "运动教练", "瑜伽达人", "游泳健将",
            "舞蹈爱好者", "篮球达人", "足球爱好者", "网球高手", "羽毛球爱好者"
        ]
        
        self.realistic_locations = [
            "北京市朝阳区", "上海市浦东新区", "广州市天河区", "深圳市南山区", "杭州市西湖区",
            "成都市锦江区", "重庆市渝中区", "南京市鼓楼区", "武汉市武昌区", "西安市碑林区",
            "苏州市姑苏区", "厦门市思明区", "青岛市市南区", "大连市中山区", "长沙市岳麓区",
            "天津市和平区", "沈阳市沈河区", "济南市历下区", "郑州市金水区", "福州市鼓楼区"
        ]
        
        self.realistic_activities = [
            {
                "type": "run",
                "title_templates": [
                    "跑步{}公里，用时{}分钟",
                    "晨跑{}公里，感觉不错",
                    "夜跑{}公里，放松心情",
                    "午间慢跑{}公里",
                    "周末长跑{}公里，挑战自我"
                ],
                "description_templates": [
                    "今天天气很好，心情愉快，跑步节奏稳定。",
                    "今天有点累，但坚持下来了，继续保持！",
                    "尝试了新路线，风景很不错，值得推荐。",
                    "今天配速提高了，感觉进步不少。",
                    "和跑友一起跑步，互相鼓励，坚持到了终点。"
                ],
                "distance_range": (2.0, 15.0),
                "duration_per_km": (4.5, 7.0),  # 分钟/公里
                "calorie_per_km": (60, 80)  # 卡路里/公里
            },
            {
                "type": "bike",
                "title_templates": [
                    "骑行{}公里，用时{}分钟",
                    "周末骑行{}公里，探索新路线",
                    "公路骑行{}公里，速度提升",
                    "山地骑行{}公里，挑战自我",
                    "通勤骑行{}公里，低碳出行"
                ],
                "description_templates": [
                    "今天骑行路线很平坦，速度很快，心情愉悦。",
                    "爬坡骑行，很有挑战性，但风景绝佳。",
                    "和车友一起骑行，互相鼓励，完成了长距离挑战。",
                    "今天尝试了新的骑行姿势，感觉效率提升了。",
                    "雨后骑行，空气清新，路面有点湿滑，需要小心。"
                ],
                "distance_range": (10.0, 60.0),
                "duration_per_km": (2.0, 4.0),  # 分钟/公里
                "calorie_per_km": (25, 40)  # 卡路里/公里
            },
            {
                "type": "walk",
                "title_templates": [
                    "健走{}公里，锻炼身体",
                    "午休散步{}公里，放松心情",
                    "傍晚漫步{}公里，欣赏美景",
                    "快走{}公里，燃脂减肥",
                    "户外健走{}公里，呼吸新鲜空气"
                ],
                "description_templates": [
                    "每天坚持走路，让身体更健康。",
                    "边走边欣赏沿途风景，心情放松。",
                    "快走锻炼，感觉很有活力。",
                    "雨后漫步，空气清新，很舒适。",
                    "和朋友一起健走，边走边聊，很愉快。"
                ],
                "distance_range": (1.0, 8.0),
                "duration_per_km": (10.0, 15.0),  # 分钟/公里
                "calorie_per_km": (40, 60)  # 卡路里/公里
            },
            {
                "type": "swim",
                "title_templates": [
                    "游泳{}米，畅快淋漓",
                    "自由泳{}米，提升技巧",
                    "蛙泳训练{}米，持续进步",
                    "泳池锻炼{}米，保持状态",
                    "游泳{}米，放松身心"
                ],
                "description_templates": [
                    "今天游泳感觉很好，动作更协调了。",
                    "尝试了新的呼吸节奏，感觉很有效。",
                    "坚持每周游泳，身体素质明显提升。",
                    "游泳后全身放松，睡眠质量也提高了。",
                    "和教练一起训练，纠正了一些动作问题。"
                ],
                "distance_range": (500, 3000),
                "duration_per_100m": (1.5, 3.0),  # 分钟/100米
                "calorie_per_100m": (10, 15)  # 卡路里/100米
            }
        ]
        
        self.realistic_routes = [
            {
                "title_templates": [
                    "{}公园跑步路线",
                    "{}河畔慢跑道",
                    "{}周边5公里经典跑道",
                    "{}健身步道",
                    "{}山间越野跑道"
                ],
                "description_templates": [
                    "这是一条非常适合晨跑的路线，路面平整，人流量小，空气清新。",
                    "这条路线沿河而建，风景优美，跑步的同时可以欣赏河景。",
                    "这是一条具有一定坡度的路线，适合想要提升体能的跑友。",
                    "这条路线环绕公园一周，全程有树荫，夏天跑步也不会太晒。",
                    "这是一条结合了道路和越野的混合路线，有一定的挑战性。"
                ],
                "distance_range": (3.0, 10.0),
                "location_landmarks": [
                    "奥林匹克公园", "世纪公园", "滨江公园", "莲花山公园", "紫竹公园",
                    "珠江", "钱塘江", "松花江", "汉江", "海河",
                    "健身步道", "城市绿道", "森林公园", "湿地公园", "山地公园"
                ]
            },
            {
                "title_templates": [
                    "{}自行车专用道",
                    "{}环线骑行路线",
                    "{}至{}经典骑行路线",
                    "{}山地车道",
                    "{}风景骑行路线"
                ],
                "description_templates": [
                    "这是一条平坦的骑行专用道，适合休闲骑行和新手练习。",
                    "这条环线路线风景优美，骑行过程中可以欣赏到不同的景色。",
                    "这是一条城市间的骑行路线，沿途有多个休息点和补给站。",
                    "这条山地车道有一定的技术难度，适合有经验的骑行者挑战。",
                    "这是一条风景优美的骑行路线，非常适合周末骑行和拍照。"
                ],
                "distance_range": (15.0, 80.0),
                "location_landmarks": [
                    "自行车专用道", "绿道", "滨海公路", "山间公路", "环湖公路",
                    "城市绿道", "森林公园", "风景区", "乡村小道", "城市快速路"
                ]
            }
        ]
        
        # 固定数量的用户ID，用于产生更真实的用户-活动关系
        self.user_ids = [f"user_{i}" for i in range(1000, 3000)]
        
    def init_database(self):
        """初始化数据库连接和表结构"""
        try:
            self.conn = pymysql.connect(**DB_CONFIG)
            self.cursor = self.conn.cursor()
            
            # 设置连接参数优化批量插入
            self.cursor.execute("SET autocommit=0")  # 关闭自动提交
            self.cursor.execute("SET unique_checks=0")  # 关闭唯一性检查
            self.cursor.execute("SET foreign_key_checks=0")  # 关闭外键检查
            
            # 创建表结构
            self._create_tables()
            
            logger.info("数据库初始化成功")
            return True
        except Exception as e:
            logger.error(f"数据库初始化失败: {e}")
            if self.conn:
                self.conn.rollback()
            return False
    
    def _create_tables(self):
        """创建数据表"""
        # 用户表
        self.cursor.execute("""
        CREATE TABLE IF NOT EXISTS `users` (
            `user_id` VARCHAR(50) PRIMARY KEY,
            `nickname` VARCHAR(100),
            `avatar_url` VARCHAR(255),
            `gender` VARCHAR(10),
            `location` VARCHAR(100),
            `follower_count` INT,
            `following_count` INT,
            `post_count` INT,
            `created_at` DATETIME
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """)
        
        # 运动记录表 - 增加route_data字段长度
        self.cursor.execute("""
        CREATE TABLE IF NOT EXISTS `activities` (
            `activity_id` VARCHAR(50) PRIMARY KEY,
            `user_id` VARCHAR(50),
            `type` VARCHAR(50),
            `title` VARCHAR(255),
            `description` TEXT,
            `distance` FLOAT,
            `duration` INT,
            `calorie` FLOAT,
            `start_time` DATETIME,
            `likes_count` INT,
            `comments_count` INT,
            `route_data` MEDIUMTEXT,
            `created_at` DATETIME,
            FOREIGN KEY (`user_id`) REFERENCES `users`(`user_id`) ON DELETE CASCADE
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """)
        
        # 运动路线表
        self.cursor.execute("""
        CREATE TABLE IF NOT EXISTS `routes` (
            `route_id` VARCHAR(50) PRIMARY KEY,
            `title` VARCHAR(255),
            `description` TEXT,
            `distance` FLOAT,
            `author_id` VARCHAR(50),
            `likes_count` INT,
            `used_count` INT,
            `location` VARCHAR(255),
            `created_at` DATETIME,
            FOREIGN KEY (`author_id`) REFERENCES `users`(`user_id`) ON DELETE CASCADE
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """)
        
        self.conn.commit()
    
    def _generate_mock_user(self, user_id):
        """生成模拟用户数据"""
        # 提取用户ID中的数字部分
        id_num = int(user_id.split('_')[-1]) if '_' in user_id else random.randint(1000, 9999)
        
        # 生成更真实的昵称
        user_name_prefix = self.realistic_user_names[id_num % len(self.realistic_user_names)]
        nickname = f"{user_name_prefix}{id_num % 1000}"
        
        # 根据用户ID确定性地生成一些属性，使同一用户的数据保持一致
        random.seed(id_num)
        gender = "male" if id_num % 2 == 0 else "female"
        location = self.realistic_locations[id_num % len(self.realistic_locations)]
        
        # 生成合理的社交数据
        follower_base = int(random.paretovariate(1.5)) % 500  # 基础粉丝数
        followers = max(follower_base, 5)  # 至少有5个粉丝
        following = int(followers * random.uniform(0.5, 1.5))  # 关注数与粉丝数相关
        post_count = int(followers * random.uniform(0.2, 1.0))  # 发帖数与粉丝数相关
        
        # 重置随机种子
        random.seed()
        
        return {
            'id': user_id,
            'nickname': nickname,
            'avatar': f"https://api.codoon.com/avatar/{user_id}.jpg",
            'gender': gender,
            'location': location,
            'follower_count': followers,
            'following_count': following,
            'post_count': post_count
        }
    
    def _generate_mock_activity(self, user_id=None, timestamp=None):
        """生成模拟活动数据"""
        # 生成ID
        activity_id = f"activity_{int(time.time() if timestamp is None else timestamp)}_{random.randint(1000, 9999)}"
        
        # 如果没有提供用户ID，则从预设用户中随机选择
        if user_id is None:
            user_id = random.choice(self.user_ids)
        
        # 随机选择活动类型
        activity_type_data = random.choice(self.realistic_activities)
        activity_type = activity_type_data["type"]
        
        # 生成合理的活动距离
        if activity_type == "swim":
            distance = random.randint(*activity_type_data["distance_range"]) / 1000  # 转换为公里
            duration_per_unit = random.uniform(*activity_type_data["duration_per_100m"]) 
            duration = int(distance * 1000 / 100 * duration_per_unit * 60)  # 转换为秒
            calorie_per_unit = random.uniform(*activity_type_data["calorie_per_100m"])
            calorie = distance * 1000 / 100 * calorie_per_unit
        else:
            distance = round(random.uniform(*activity_type_data["distance_range"]), 2)
            duration_per_unit = random.uniform(*activity_type_data["duration_per_km"])
            duration = int(distance * duration_per_unit * 60)  # 转换为秒
            calorie_per_unit = random.uniform(*activity_type_data["calorie_per_km"])
            calorie = distance * calorie_per_unit
        
        # 生成活动标题
        title_template = random.choice(activity_type_data["title_templates"])
        if activity_type == "swim":
            title = title_template.format(int(distance * 1000))
        else:
            title = title_template.format(distance, int(duration/60))
        
        # 生成活动描述
        description = random.choice(activity_type_data["description_templates"])
        
        # 生成合理的时间
        if timestamp is None:
            current_time = datetime.now()
            days_ago = random.randint(1, 365)  # 增加时间范围到1年
            hours = random.randint(6, 21)  # 合理的运动时间
            start_time = current_time - timedelta(days=days_ago)
            start_time = start_time.replace(hour=hours, minute=random.randint(0, 59))
            timestamp = int(start_time.timestamp())
        else:
            start_time = datetime.fromtimestamp(timestamp)
        
        # 生成点赞和评论数 - 基于帕累托分布，更符合社交平台数据分布特征
        likes_count = int(random.paretovariate(1.5)) % 150
        comments_count = int(likes_count * random.uniform(0.1, 0.5))
        
        # 简化路线数据，仅保留必要字段以减少数据量
        route_data = {
            "start_point": {
                "lat": 39.90 + random.uniform(-0.2, 0.2),
                "lng": 116.40 + random.uniform(-0.2, 0.2)
            },
            "end_point": {
                "lat": 39.90 + random.uniform(-0.2, 0.2),
                "lng": 116.40 + random.uniform(-0.2, 0.2)
            },
            # 只保存5个路径点，足够展示和减少存储量
            "path_points": [{
                "lat": 39.90 + random.uniform(-0.2, 0.2),
                "lng": 116.40 + random.uniform(-0.2, 0.2),
                "timestamp": timestamp + int(duration * i / 5),
                "altitude": 50 + random.uniform(-10, 20),
                "speed": random.uniform(2.0, 30.0)
            } for i in range(5)]
        }
        
        return {
            'id': activity_id,
            'user_id': user_id,
            'type': activity_type,
            'title': title,
            'description': description,
            'distance': distance,
            'duration': duration,
            'calorie': round(calorie, 2),
            'start_time': timestamp,
            'likes_count': likes_count,
            'comments_count': comments_count,
            'route_data': json.dumps(route_data)  # 直接序列化为JSON字符串
        }
    
    def _generate_mock_route(self, author_id=None, timestamp=None):
        """生成模拟路线数据"""
        # 生成ID
        route_id = f"route_{int(time.time() if timestamp is None else timestamp)}_{random.randint(1000, 9999)}"
        
        # 如果没有提供作者ID，则从预设用户中随机选择
        if author_id is None:
            author_id = random.choice(self.user_ids)
        
        # 随机选择路线类型
        route_type_data = random.choice(self.realistic_routes)
        
        # 生成合理的路线距离
        distance = round(random.uniform(*route_type_data["distance_range"]), 2)
        
        # 生成路线标题
        location_landmark = random.choice(route_type_data["location_landmarks"])
        title_template = random.choice(route_type_data["title_templates"])
        
        # 处理不同的标题模板情况
        if "{}至{}" in title_template:
            location_landmark2 = random.choice(route_type_data["location_landmarks"])
            while location_landmark2 == location_landmark:
                location_landmark2 = random.choice(route_type_data["location_landmarks"])
            title = title_template.format(location_landmark, location_landmark2)
        else:
            title = title_template.format(location_landmark)
        
        # 生成路线描述
        description = random.choice(route_type_data["description_templates"])
        
        # 生成合理的点赞和使用数 - 使用帕累托分布
        likes_count = int(random.paretovariate(1.2)) % 300
        used_count = int(likes_count * random.uniform(0.5, 4.0))
        
        # 随机选择地点
        location = random.choice(self.realistic_locations)
        
        # 生成时间
        if timestamp is None:
            current_time = datetime.now()
            days_ago = random.randint(1, 365)  # 增加时间范围到1年
            created_at = current_time - timedelta(days=days_ago)
        else:
            created_at = datetime.fromtimestamp(timestamp)
        
        return {
            'route_id': route_id,
            'author_id': author_id,
            'title': title,
            'description': description,
            'distance': distance,
            'likes_count': likes_count,
            'used_count': used_count,
            'location': location,
            'created_at': created_at
        }
    
    def generate_and_save_users(self, count=1000):
        """生成并批量保存用户数据"""
        try:
            logger.info(f"开始生成{count}条用户数据...")
            
            # 预生成所有用户数据
            users_data = []
            user_batch = []
            now = datetime.now()
            
            # 使用tqdm显示进度条
            for i in tqdm(range(count)):
                user_id = self.user_ids[i % len(self.user_ids)]
                user_data = self._generate_mock_user(user_id)
                
                # 构建SQL值元组
                user_tuple = (
                    user_data['id'],
                    user_data['nickname'],
                    user_data['avatar'],
                    user_data['gender'],
                    user_data['location'],
                    user_data['follower_count'],
                    user_data['following_count'],
                    user_data['post_count'],
                    now
                )
                
                user_batch.append(user_tuple)
                users_data.append(user_data)
                
                # 达到批处理大小后执行批量插入
                if len(user_batch) >= self.batch_size:
                    self._batch_insert_users(user_batch)
                    user_batch = []
            
            # 处理剩余的用户数据
            if user_batch:
                self._batch_insert_users(user_batch)
            
            self.conn.commit()
            logger.info(f"成功生成并保存{count}条用户数据")
            return users_data
        except Exception as e:
            logger.error(f"生成用户数据时出错: {e}")
            self.conn.rollback()
            return []
    
    def _batch_insert_users(self, user_batch):
        """批量插入用户数据"""
        sql = """
        INSERT IGNORE INTO `users` 
        (`user_id`, `nickname`, `avatar_url`, `gender`, `location`,
         `follower_count`, `following_count`, `post_count`, `created_at`)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        self.cursor.executemany(sql, user_batch)
    
    def generate_and_save_activities(self, count=10000):
        """生成并批量保存活动数据"""
        try:
            logger.info(f"开始生成{count}条活动数据...")
            
            # 获取所有已保存的用户ID
            self.cursor.execute("SELECT user_id FROM users")
            user_ids = [row[0] for row in self.cursor.fetchall()]
            
            if not user_ids:
                logger.error("没有找到用户数据，无法生成活动")
                return []
            
            # 预生成所有活动数据
            activities_data = []
            activity_batch = []
            now = datetime.now()
            
            # 使用tqdm显示进度条
            for i in tqdm(range(count)):
                # 为每个活动随机选择一个用户
                user_id = random.choice(user_ids)
                timestamp = int((now - timedelta(days=random.randint(1, 365))).timestamp())
                activity_data = self._generate_mock_activity(user_id, timestamp)
                
                # 构建SQL值元组
                activity_tuple = (
                    activity_data['id'],
                    activity_data['user_id'],
                    activity_data['type'],
                    activity_data['title'],
                    activity_data['description'],
                    activity_data['distance'],
                    activity_data['duration'],
                    activity_data['calorie'],
                    datetime.fromtimestamp(activity_data['start_time']),
                    activity_data['likes_count'],
                    activity_data['comments_count'],
                    activity_data['route_data'],
                    now
                )
                
                activity_batch.append(activity_tuple)
                activities_data.append(activity_data)
                
                # 达到批处理大小后执行批量插入
                if len(activity_batch) >= self.batch_size:
                    self._batch_insert_activities(activity_batch)
                    activity_batch = []
            
            # 处理剩余的活动数据
            if activity_batch:
                self._batch_insert_activities(activity_batch)
            
            self.conn.commit()
            logger.info(f"成功生成并保存{count}条活动数据")
            return activities_data
        except Exception as e:
            logger.error(f"生成活动数据时出错: {e}")
            self.conn.rollback()
            return []
    
    def _batch_insert_activities(self, activity_batch):
        """批量插入活动数据"""
        sql = """
        INSERT IGNORE INTO `activities` 
        (`activity_id`, `user_id`, `type`, `title`, `description`, `distance`,
         `duration`, `calorie`, `start_time`, `likes_count`, `comments_count`,
         `route_data`, `created_at`)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        self.cursor.executemany(sql, activity_batch)
    
    def generate_and_save_routes(self, count=5000):
        """生成并批量保存路线数据"""
        try:
            logger.info(f"开始生成{count}条路线数据...")
            
            # 获取所有已保存的用户ID
            self.cursor.execute("SELECT user_id FROM users")
            user_ids = [row[0] for row in self.cursor.fetchall()]
            
            if not user_ids:
                logger.error("没有找到用户数据，无法生成路线")
                return []
            
            # 预生成所有路线数据
            routes_data = []
            route_batch = []
            now = datetime.now()
            
            # 使用tqdm显示进度条
            for i in tqdm(range(count)):
                # 为每个路线随机选择一个作者
                author_id = random.choice(user_ids)
                timestamp = int((now - timedelta(days=random.randint(1, 365))).timestamp())
                route_data = self._generate_mock_route(author_id, timestamp)
                
                # 构建SQL值元组
                route_tuple = (
                    route_data['route_id'],
                    route_data['title'],
                    route_data['description'],
                    route_data['distance'],
                    route_data['author_id'],
                    route_data['likes_count'],
                    route_data['used_count'],
                    route_data['location'],
                    route_data['created_at']
                )
                
                route_batch.append(route_tuple)
                routes_data.append(route_data)
                
                # 达到批处理大小后执行批量插入
                if len(route_batch) >= self.batch_size:
                    self._batch_insert_routes(route_batch)
                    route_batch = []
            
            # 处理剩余的路线数据
            if route_batch:
                self._batch_insert_routes(route_batch)
            
            self.conn.commit()
            logger.info(f"成功生成并保存{count}条路线数据")
            return routes_data
        except Exception as e:
            logger.error(f"生成路线数据时出错: {e}")
            self.conn.rollback()
            return []
    
    def _batch_insert_routes(self, route_batch):
        """批量插入路线数据"""
        sql = """
        INSERT IGNORE INTO `routes` 
        (`route_id`, `title`, `description`, `distance`, `author_id`,
         `likes_count`, `used_count`, `location`, `created_at`)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        self.cursor.executemany(sql, route_batch)
    
    def generate_all_data(self, users_count=2000, activities_count=15000, routes_count=5000):
        """生成所有数据"""
        try:
            # 重新创建数据库
            self.cursor.execute("DROP DATABASE IF EXISTS codoon_data")
            self.cursor.execute("CREATE DATABASE codoon_data CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
            self.cursor.execute("USE codoon_data")
            
            # 创建表结构
            self._create_tables()
            
            # 生成用户数据
            start_time = time.time()
            self.generate_and_save_users(users_count)
            logger.info(f"生成用户数据耗时: {time.time() - start_time:.2f}秒")
            
            # 生成活动数据
            start_time = time.time()
            self.generate_and_save_activities(activities_count)
            logger.info(f"生成活动数据耗时: {time.time() - start_time:.2f}秒")
            
            # 生成路线数据
            start_time = time.time()
            self.generate_and_save_routes(routes_count)
            logger.info(f"生成路线数据耗时: {time.time() - start_time:.2f}秒")
            
            # 恢复数据库设置
            self.cursor.execute("SET unique_checks=1")
            self.cursor.execute("SET foreign_key_checks=1")
            self.cursor.execute("SET autocommit=1")
            
            logger.info(f"成功生成总计{users_count+activities_count+routes_count}条数据!")
        except Exception as e:
            logger.error(f"生成数据过程中出错: {e}")
            self.conn.rollback()
        finally:
            # 关闭数据库连接
            if self.cursor:
                self.cursor.close()
            if self.conn:
                self.conn.close()

if __name__ == "__main__":
    # 生成2万条数据 (2000用户 + 15000活动 + 5000路线)
    total_start_time = time.time()
    generator = CodoonBatchGenerator(batch_size=1000)
    generator.init_database()
    generator.generate_all_data(2000, 15000, 5000)
    logger.info(f"总耗时: {time.time() - total_start_time:.2f}秒") 